import pandas as pd
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
import SimpleITK as sitk
import random


def pprint(sitkImage):
    print("原点位置:{}".format(sitkImage.GetOrigin()))
    print("尺寸：{}".format(sitkImage.GetSize()))
    print("体素大小(x,y,z):{}".format(sitkImage.GetSpacing()))
    print("图像方向:{}".format(sitkImage.GetDirection()))

    # 查看图像相关的纬度信息
    print("维度:{}".format(sitkImage.GetDimension()))
    print("宽度:{}".format(sitkImage.GetWidth()))
    print("高度:{}".format(sitkImage.GetHeight()))
    print("深度(层数):{}".format(sitkImage.GetDepth()))


def sitk_read_raw(img_path, mode='train'):
    # 读取nifit格式，设置窗口窗位，并缩放到0到255之间，并转换成numpy格式
    sitkImage = sitk.ReadImage(img_path)
    intensityWindowingFilter = sitk.IntensityWindowingImageFilter()
    # 转换成0到255之间
    # intensityWindowingFilter.SetOutputMinimum(0)
    # intensityWindowingFilter.SetOutputMaximum(255)
    if 'mask' not in img_path:
        # 调窗宽窗位
        # 相当于数据增强中的 改变对比度
        if mode == 'train':
            max_num = random.randint(1000, 1900)
            min_num = random.randint(-600, 0)
            intensityWindowingFilter.SetWindowMaximum(max_num)
            intensityWindowingFilter.SetWindowMinimum(min_num)
        else:
            intensityWindowingFilter.SetWindowMaximum(1500)
            intensityWindowingFilter.SetWindowMinimum(-600)
    # 设置窗宽窗位，并进行缩放
    sitkImage = intensityWindowingFilter.Execute(sitkImage)
    nda = sitk.GetArrayFromImage(sitkImage)
    return np.transpose(nda, (1, 2, 0))


def Scale(img_path):
    # 通过RescaleIntensity方法既可以进行缩放，没有窗位的调整
    sitkImage = sitk.ReadImage(img_path)
    sitkImage = sitk.RescaleIntensity(sitkImage)
    nda = sitk.GetArrayFromImage(sitkImage)
    return np.transpose(nda, (1, 2, 0))


def read_nii(filepath):
    '''
    Reads .nii file and returns pixel array
    '''
    ct_scan = nib.load(filepath)
    array = ct_scan.get_fdata()
    array = np.rot90(np.array(array))
    print(array.shape)
    return array


if __name__ == '__main__':
    root = 'E:\\2022Project\\data\\COVID-19-CT'

    raw_data = pd.read_csv(root + '/metadata.csv')
    for cc in raw_data.columns:
        raw_data[cc] = raw_data[cc].apply(lambda x: root + '/' + '/'.join(x.split('/')[-2:]))

    for name in ['lung_mask','infection_mask'][:1]:
        print('processing {}'.format(name))
        category = 1
        sample_ct = None
        sample_gt = None
        ct_scan = raw_data['ct_scan'].tolist()
        mask = raw_data[name].tolist()
        sitkImage = sitk.ReadImage(ct_scan[1])
        z = int(sitkImage.GetSize()[2] / 2)
        slice = sitk.GetArrayFromImage(sitkImage)[z, :, :]
        print("调节窗口窗位之前CT值的范围位为{}~{}".format(np.min(slice), np.max(slice)))
        plt.figure(figsize=(5, 5))
        plt.imshow(slice, 'gray')
        plt.show()

        sitkimage = sitk_read_raw(ct_scan[1],'valid')
        z = int(sitkimage.shape[2] / 2)
        slice = sitkimage[:, :, z]
        print("调节窗口窗位之后CT值的范围位为{}~{}".format(np.min(slice), np.max(slice)))
        plt.figure(figsize=(5, 5))
        plt.imshow(slice, 'gray')
        plt.show()

        sitkimage = sitk_read_raw(mask[1])
        z = int(sitkimage.shape[2] / 2)
        slice = sitkimage[:, :, z]
        print("调节窗口窗位之后CT值的范围位为{}~{}".format(np.min(slice), np.max(slice)))
        plt.figure(figsize=(5, 5))
        plt.imshow(slice, 'gray')
        plt.show()


